{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy.optimize import fsolve\n",
    "from math import pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# esthetic definitions for the plots\n",
    "font = {'size'   : 16, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'\n",
    "plt.rcParams[\"figure.figsize\"] = [8.0,6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################\n",
    "#\n",
    "# Varying parameter (others fixed to default)\n",
    "#\n",
    "# With the input suntax of class <= 2.9 we used: annihilation = 1.e-5 m^3/s/Kg\n",
    "# With the new syntax this is equivalent to DM_annihilation_efficiency = 1.11e-22 m^3/s/J\n",
    "# (the ratio is a factor (c/[1 m/s])**2 = 9.e16)\n",
    "#\n",
    "var_name = 'DM_annihilation_efficiency'\n",
    "var_array = np.linspace(0,1.11e-22,5)\n",
    "var_num = len(var_array)\n",
    "var_legend = r'$p_\\mathrm{ann}$'\n",
    "var_figname = 'pann'\n",
    "#\n",
    "#############################################\n",
    "#\n",
    "# Fixed settings\n",
    "#\n",
    "common_settings = {# LambdaCDM parameters\n",
    "                   'h':0.67556,\n",
    "                   'omega_b':0.022032,\n",
    "                   'omega_cdm':0.12038,\n",
    "                   'A_s':2.215e-9,\n",
    "                   'n_s':0.9619,\n",
    "                   'tau_reio':0.0925,\n",
    "                   # output and precision parameters\n",
    "                   'output':'tCl,pCl,lCl,mPk',\n",
    "                   'lensing':'yes',\n",
    "                   'P_k_max_1/Mpc':3.0,\n",
    "                   'l_switch_limber':9\n",
    "                   }\n",
    "#\n",
    "# arrays for output\n",
    "#\n",
    "kvec = np.logspace(-4,np.log10(3),1000)\n",
    "legarray = []\n",
    "twopi = 2.*pi\n",
    "#\n",
    "# Create figures\n",
    "#\n",
    "fig_Pk, ax_Pk = plt.subplots()\n",
    "fig_TT, ax_TT = plt.subplots()\n",
    "fig_EE, ax_EE = plt.subplots()\n",
    "fig_PP, ax_PP = plt.subplots()\n",
    "#\n",
    "M = Class()\n",
    "#\n",
    "# loop over varying parameter values\n",
    "#\n",
    "for i,var in enumerate(var_array):\n",
    "    #\n",
    "    print (' * Compute with %s=%e'%(var_name,var))\n",
    "    #\n",
    "    # deal with colors and legends\n",
    "    #\n",
    "    if i == 0:\n",
    "        var_color = 'k'\n",
    "        var_alpha = 1.\n",
    "        legarray.append(r'ref. $\\Lambda CDM$')\n",
    "    else:\n",
    "        var_color = 'r'\n",
    "        var_alpha = 1.*i/(var_num-1.)\n",
    "    if i == var_num-1:\n",
    "        legarray.append(var_legend)  \n",
    "    #    \n",
    "    # call CLASS\n",
    "    #\n",
    "    M.set(common_settings)\n",
    "    M.set({var_name:var})\n",
    "    M.compute()\n",
    "    #\n",
    "    # get Cls\n",
    "    #\n",
    "    clM = M.lensed_cl(2500)\n",
    "    ll = clM['ell'][2:]\n",
    "    clTT = clM['tt'][2:]\n",
    "    clEE = clM['ee'][2:]\n",
    "    clPP = clM['pp'][2:]\n",
    "    #\n",
    "    # get P(k) for common k values\n",
    "    #\n",
    "    pkM = []\n",
    "    for k in kvec:\n",
    "        pkM.append(M.pk(k,0.))\n",
    "    #    \n",
    "    # plot P(k)\n",
    "    #\n",
    "    ax_Pk.loglog(kvec,np.array(pkM),color=var_color,alpha=var_alpha,linestyle='-')\n",
    "    #\n",
    "    # plot C_l^TT\n",
    "    #\n",
    "    ax_TT.semilogx(ll,clTT*ll*(ll+1)/twopi,color=var_color,alpha=var_alpha,linestyle='-')\n",
    "    #\n",
    "    # plot Cl EE \n",
    "    #\n",
    "    ax_EE.loglog(ll,clEE*ll*(ll+1)/twopi,color=var_color,alpha=var_alpha,linestyle='-')\n",
    "    #\n",
    "    # plot Cl phiphi\n",
    "    #\n",
    "    ax_PP.loglog(ll,clPP*ll*(ll+1)*ll*(ll+1)/twopi,color=var_color,alpha=var_alpha,linestyle='-')\n",
    "    #\n",
    "    # reset CLASS\n",
    "    #\n",
    "    M.empty()    \n",
    "#\n",
    "# output of P(k) figure\n",
    "#\n",
    "ax_Pk.set_xlim([1.e-4,3.])\n",
    "ax_Pk.set_xlabel(r'$k \\,\\,\\,\\, [h/\\mathrm{Mpc}]$')\n",
    "ax_Pk.set_ylabel(r'$P(k) \\,\\,\\,\\, [\\mathrm{Mpc}/h]^3$')\n",
    "ax_Pk.legend(legarray)\n",
    "fig_Pk.tight_layout()\n",
    "fig_Pk.savefig('varying_%s_Pk.pdf' % var_figname)\n",
    "#\n",
    "# output of C_l^TT figure\n",
    "#      \n",
    "ax_TT.set_xlim([2,2500])\n",
    "ax_TT.set_xlabel(r'$\\ell$')\n",
    "ax_TT.set_ylabel(r'$[\\ell(\\ell+1)/2\\pi]  C_\\ell^\\mathrm{TT}$')\n",
    "ax_TT.legend(legarray)\n",
    "fig_TT.tight_layout()\n",
    "fig_TT.savefig('varying_%s_cltt.pdf' % var_figname)\n",
    "#\n",
    "# output of C_l^EE figure\n",
    "#    \n",
    "ax_EE.set_xlim([2,2500])\n",
    "ax_EE.set_xlabel(r'$\\ell$')\n",
    "ax_EE.set_ylabel(r'$[\\ell(\\ell+1)/2\\pi]  C_\\ell^\\mathrm{EE}$')\n",
    "ax_EE.legend(legarray)\n",
    "fig_EE.tight_layout()\n",
    "fig_EE.savefig('varying_%s_clee.pdf' % var_figname)\n",
    "#\n",
    "# output of C_l^pp figure\n",
    "#   \n",
    "ax_PP.set_xlim([10,2500])\n",
    "ax_PP.set_xlabel(r'$\\ell$')\n",
    "ax_PP.set_ylabel(r'$[\\ell^2(\\ell+1)^2/2\\pi]  C_\\ell^\\mathrm{\\phi \\phi}$')\n",
    "ax_PP.legend(legarray)\n",
    "fig_PP.tight_layout()\n",
    "fig_PP.savefig('varying_%s_clpp.pdf' % var_figname)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
